#include "header.h"

static ptrdiff_t	i, j, k, n0, n1, n2, n;
static fftw_plan 	plan, iplan, grvft, tmpb; //focon; // incon;
static int		myid, wisdom_flag, np, heat_flag, gFlag, mFlag;
static double		overN, c, l, dx, dt, multip, Nsq, Npart, *x, *y, *z, rr;
static double		*dealias, kcut, kwid;

void init_ffts(ctrl_ptr ptr, work_ptr work)
{ 
  MPI_Comm_rank(COMM, &myid);
  n  = ptr->n; gFlag = ptr->gFlag; mFlag = ptr->mFlag;
  n0 = ptr->n; n1 = ptr->n; n2 = ptr->n;
  heat_flag = ptr->heat; l = ptr->l;
  dx = l/n; dt = (ptr->cfl)*dx*dx;
  overN = 1./(n*n*n); Npart = ptr->Npart;
  x = malloc(n*sizeof(double));
  y = malloc(n*sizeof(double));
  z = malloc(n*sizeof(double));
  MPI_Barrier(COMM);  

  if (myid == 0) wisdom_flag = fftw_import_wisdom_from_filename("wisdom.fft");
  MPI_Bcast(&wisdom_flag, sizeof(int), MPI_BYTE, 0, COMM);
  MPI_Barrier(COMM);
  if (wisdom_flag == 1) {
    fftw_mpi_broadcast_wisdom(COMM);
    MPI_Barrier(COMM);  
    plan  = fftw_mpi_plan_dft_3d(n0,n1,n2,work->ft,work->ft,COMM,FFTW_FORWARD,FMODE);
    iplan = fftw_mpi_plan_dft_3d(n0,n1,n2,work->ft,work->ft,COMM,FFTW_BACKWARD,FMODE);
    grvft = fftw_mpi_plan_dft_3d(n0,n1,n2,work->grv,work->grv,COMM,FFTW_FORWARD,FMODE);
    tmpb = fftw_mpi_plan_dft_3d(n0,n1,n2,work->tmp,work->tmp,COMM,FFTW_BACKWARD,FMODE); 
    MPI_Barrier(COMM);  
    if (myid == 0) printf("Loaded wisdom\n");
  } else {
    if (myid == 0) printf("Creating new plans\n");
    MPI_Barrier(COMM);  
    plan  = fftw_mpi_plan_dft_3d(n0,n1,n2,work->ft,work->ft,COMM,FFTW_FORWARD,FMODE);
    iplan = fftw_mpi_plan_dft_3d(n0,n1,n2,work->ft,work->ft,COMM,FFTW_BACKWARD,FMODE);
    grvft = fftw_mpi_plan_dft_3d(n0,n1,n2,work->grv,work->grv,COMM,FFTW_FORWARD,FMODE);
    tmpb = fftw_mpi_plan_dft_3d(n0,n1,n2,work->tmp,work->tmp,COMM,FFTW_BACKWARD,FMODE); 
    MPI_Barrier(COMM);
    fftw_mpi_gather_wisdom(COMM);
    if (myid == 0) {
      fftw_export_wisdom_to_filename("wisdom.fft");
      printf("Wisdom saved successfully\n");
    }
  } 
  MPI_Barrier(COMM);
  multip = l*l*l*overN*overN;
}

void ft_potential()
{
  fftw_execute(grvft);
}
void fft_laplace(work_ptr in)
{
  int 		kx, ky, kz;
  double 	DD;

  c = 2*PI/l;
  for (i = 0; i < (in->alloc_local); i++) in->ft[i] = in->wrk[i];
  MPI_Barrier(COMM);  
  fftw_execute(plan);
  MPI_Barrier(COMM);  
  for (i = 0; i < (in->l_n0); i++) for (j = 0; j < n; j++) for (k = 0; k < n; k++)
  {
    kx = i + myid*(in->l_n0);
    ky = j;
    kz = k;
    if (kx > n/2) kx = kx - n;
    if (ky > n/2) ky = ky - n;
    if (kz > n/2) kz = kz - n;
    DD = c*c*(kx*kx + ky*ky + kz*kz);
    in->ft[k+n*(j+n*i)] = -DD*(in->ft[k+n*(j+n*i)])*overN;      
  }
  MPI_Barrier(COMM);  
  fftw_execute(iplan);
  MPI_Barrier(COMM);  
}
void lin_step(work_ptr in, fftw_complex step)
{
  int 		kx, ky, kz;
  double 	DD;

  c = 2*PI/l;
  MPI_Barrier(COMM);  
  for (i = 0; i < (in->alloc_local); i++) in->ft[i] = in->wrk[i];
  MPI_Barrier(COMM);
  fftw_execute(plan);
  MPI_Barrier(COMM);  
  for (i = 0; i < (in->l_n0); i++) for (j = 0; j < n; j++) for (k = 0; k < n; k++)
  {
    kx = i + myid*(in->l_n0);
    ky = j;
    kz = k;
    if (kx > n/2) kx = kx - n;
    if (ky > n/2) ky = ky - n;
    if (kz > n/2) kz = kz - n;
    DD = c*c*(kx*kx + ky*ky + kz*kz);
    in->ft[k+n*(j+n*i)] = (in->ft[k+n*(j+n*i)])*cexp(-1.I*DD*step*dt)*overN;    
  }
  MPI_Barrier(COMM);  
  fftw_execute(iplan); 
  MPI_Barrier(COMM);
  for (i = 0; i < (in->alloc_local); i++) in->wrk[i] = in->ft[i];
  if (heat_flag == 1) 
  {
    Nsq = sqrt(particles(in)/Npart);
    for (i = 0; i < (in->alloc_local); i++) in->wrk[i] = (in->wrk[i])/Nsq;
  }
  MPI_Barrier(COMM);  
}

void convolution(work_ptr in)
{
  int 		kx, ky, kz;
  double 	x,y,z,rr;
  c = 2*PI/l;
  MPI_Barrier(COMM);  
  for (i = 0; i < (in->l_n0)*n*n; i++) in->ft[i] = pow(cabs(in->wrk[i]),2);
  /*for (i = 0; i < (in->l_n0); i++) for (j = 0; j < n; j++) for (k = 0; k < n; k++)
  {
	x = (myid*(in->l_n0)+i+0.5)*dx - l/2;
        y = (j+0.5)*dx - l/2;
	z = (k+0.5)*dx - l/2;
        rr = sqrt(x*x+y*y+z*z);
  	(in->wrk[k+n*(j+n*i)]) = 1*cexp(-1*rr*rr);
	in->ft[k+n*(j+n*i)] = pow(cabs(in->wrk[k+n*(j+n*i)]),2);
  }*/
  MPI_Barrier(COMM);
  fftw_execute(plan);
  MPI_Barrier(COMM);
  for (i = 0; i < (in->l_n0); i++) for (j = 0; j < n; j++) for (k = 0; k < n; k++)
  {
    kx = i + myid*(in->l_n0);
    ky = j;
    kz = k;
    if (kx > n/2) kx = kx - n;
    if (ky > n/2) ky = ky - n;
    if (kz > n/2) kz = kz - n;   
    if (gFlag == 1) in->ft[k+n*(j+n*i)] = multip*cabs(in->grv[k+n*(j+n*i)])*(in->ft[k+n*(j+n*i)]);
    else in->ft[k+n*(j+n*i)] = multip*(in->grv[k+n*(j+n*i)])*(in->ft[k+n*(j+n*i)])*cexp(-0.5I*l*(1 + 1./n)*c*(kx+ky+kz)); 
  }
  //printf("Here 3.0\t gFlag = %d\n", gFlag);
  MPI_Barrier(COMM);
  fftw_execute(iplan);
  MPI_Barrier(COMM);
}
void nnl_step(work_ptr in, fftw_complex step)
{
  convolution(in);
  MPI_Barrier(COMM);  
  for (i = 0; i < (in->l_n0); i++) for (j = 0; j < n; j++) for (k = 0; k < n; k++)
  {	
    in->wrk[k+n*(j+n*i)] = in->wrk[k+n*(j+n*i)]*cexp( 1.I*step*dt*(in->ft[k+n*(j+n*i)]));
  }
  MPI_Barrier(COMM);  
}

void mask(work_ptr in)
{
  for (i = 0; i < (in->l_n0); i++) {
    x[i] = (myid*(in->l_n0)+i+0.5)*dx - l/2;
    for (j = 0; j < n; j++) {
      y[j] = (j+0.5)*dx - l/2;
      for (k = 0; k < n; k++) {	
        z[k] = (k+0.5)*dx - l/2;
        rr = sqrt(x[i]*x[i]+y[j]*y[j]+z[k]*z[k]);
        //rr = sqrt(x[i]*x[i] + x[j]*x[j] + x[k]*x[k]);   
        in->wrk[k+n*(j+n*i)] = in->wrk[k+n*(j+n*i)]*(0.9 + 0.1*cexp(-pow(rr/(0.45*l),10)));
      }
    }
  }
}


double particles(work_ptr in)
{
  double npart = 0;
  double Ntot;
  
  MPI_Comm_rank(COMM, &myid);
  for (i = 0; i < (in->alloc_local); i++) in->ft[i] = in->wrk[i];
  MPI_Barrier(COMM);  
  fftw_execute(plan);
  MPI_Barrier(COMM);
  for (i = 0; i < (in->l_n0)*n*n; i++) npart = npart + cabs(in->ft[i])*cabs(in->ft[i]);
  MPI_Barrier(COMM);  
  MPI_Allreduce(&npart, &Ntot, 1, MPI_DOUBLE, MPI_SUM, COMM);
  MPI_Barrier(COMM);  

  //if ( myid == 0 ) printf("Total number is %f\n", Ntot*multip);
  return Ntot*multip;
}

void my_fft_bench(work_ptr work)
{
  MPI_Barrier(COMM);
  fftw_execute(plan);
  MPI_Barrier(COMM);  
  fftw_execute(iplan);
  MPI_Barrier(COMM);
  for (i = 0; i < (work->alloc_local); i++) work->wrk[i] = work->wrk[i]*overN;
  MPI_Barrier(COMM);
}
